# coding=utf-8
import torch
from torch import nn, optim


# 全连接神经网络
class FC(torch.nn.Module):
    def __init__(self, in_dim, fc_hidden_list):
        super(FC, self).__init__()

        self.layers = []
        self.layers.append(nn.Sequential(nn.Linear(in_dim, fc_hidden_list[0]), nn.ReLU(True)))
        for i in range(len(fc_hidden_list) - 1):
            self.layers.append(nn.Sequential(nn.Linear(fc_hidden_list[i], fc_hidden_list[i + 1]), nn.ReLU(True)))
        self.layers.append(nn.Sequential(nn.Linear(fc_hidden_list[-1], 1, nn.ReLU(True))))
        self.layers = nn.ModuleList(self.layers)

    def forward(self, x_fc):
        for layer in self.layers:
            x_fc = layer(x_fc)
        return x_fc

    @staticmethod
    def train_and_predict(x_train_fc, y_train, x_test_fc, y_test, fc_hidden_list, learning_rate, num_epochs):
        x_train_fc = torch.from_numpy(x_train_fc)
        y_train = torch.from_numpy(y_train)
        x_test_fc = torch.from_numpy(x_test_fc)
        y_test = torch.from_numpy(y_test)

        model = FC(x_train_fc.shape[1], fc_hidden_list)
        if torch.cuda.is_available():
            model = model.cuda()
            x_train_fc = x_train_fc.cuda()
            y_train = y_train.cuda()
            x_test_fc = x_test_fc.cuda()
            y_test = y_test.cuda()

        criterion = nn.MSELoss()
        optimizer = optim.Adam(model.parameters(), lr=learning_rate)
        loss_changes = []

        fc_train_out = None
        for epoch in range(num_epochs):

            fc_train_out = model.forward(x_train_fc)

            loss = criterion(fc_train_out, y_train.unsqueeze(1))

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if epoch % 100 == 0:
                print('Epoch[{}/{}], loss:{:.6f}'.format(epoch + 1, num_epochs, loss.item()))
            fc_test_out = model.forward(x_test_fc)
            loss_test = criterion(fc_test_out, y_test.unsqueeze(1))
            loss_changes.append([loss.item(), loss_test.item()])

        model.eval()

        fc_test_out = model.forward(x_test_fc)
        fc_test_out = fc_test_out.cpu().data.numpy()
        fc_train_out = fc_train_out.cpu().data.numpy()

        return fc_train_out, fc_test_out, loss_changes
